package com.jessechanrui.template.common.template.auth.config;

import com.jessechanrui.template.common.template.auth.application.dto.UserDTO;
import com.jessechanrui.template.common.template.auth.common.AuthorizationSupport;
import com.jessechanrui.template.common.template.auth.config.executor.MdcThreadPoolTaskExecutor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.aop.interceptor.AsyncUncaughtExceptionHandler;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableAsync;

import java.util.concurrent.Executor;

@Slf4j
@EnableAsync
@Configuration
public class AsyncAutoConfiguration implements AsyncConfigurer {

    @Override
    public Executor getAsyncExecutor() {
        MdcThreadPoolTaskExecutor executor = new MdcThreadPoolTaskExecutor();
        //核心线程数5：线程池创建时候初始化的线程数
        executor.setCorePoolSize(5);
        //最大线程数5：线程池最大的线程数，只有在缓冲队列满了之后才会申请超过核心线程数的线程
        executor.setMaxPoolSize(5);
        //缓冲队列500：用来缓冲执行任务的队列
        executor.setQueueCapacity(500);
        //允许线程的空闲时间60秒：当超过了核心线程出之外的线程在空闲时间到达之后会被销毁
        executor.setKeepAliveSeconds(60);
        //线程池名的前缀：设置好了之后可以方便我们定位处理任务所在的线程池
        executor.setThreadNamePrefix("Aicr-Console-Async-");
        // 设置TaskDecorator来传递上下文
        executor.setTaskDecorator(new TenantContextDecorator());
        executor.initialize();
        return executor;
    }

    @Override
    public AsyncUncaughtExceptionHandler getAsyncUncaughtExceptionHandler() {
        return (ex, method, params) -> log.error("Exception in async method: {}", method.getName(), ex);
    }

    public static class TenantContextDecorator implements TaskDecorator {
        @Override
        public Runnable decorate(Runnable runnable) {
            // 捕获主线程的令牌与租户ID
            UserDTO currentUser = AuthorizationSupport.getCurrentUser();
            String token = AuthorizationSupport.getCurrentUserToken();
            return () -> {
                try {
                    // 在异步线程中设置用户与访问令牌
                    AuthorizationSupport.setUser(token, currentUser);
                    // 执行异步任务
                    runnable.run();
                } finally {
                    // 清理上下文
                    AuthorizationSupport.removeUser();
                }
            };
        }
    }

}
